"""
数据库初始化脚本
创建数据库、表结构和初始数据
"""
import sys
from pathlib import Path

# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import text
from loguru import logger
from app.core.config import settings
from app.models.database import engine, Base, check_db_connection

logger.add("logs/init_db.log", rotation="1 day", level="INFO")


def create_database():
    """创建数据库（如果不存在）"""
    logger.info("检查数据库是否存在...")
    
    # 连接到postgres数据库
    from sqlalchemy import create_engine
    postgres_url = f"postgresql://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:{settings.DATABASE_PORT}/postgres"
    postgres_engine = create_engine(postgres_url, isolation_level="AUTOCOMMIT")
    
    try:
        with postgres_engine.connect() as conn:
            # 检查数据库是否存在
            result = conn.execute(
                text(f"SELECT 1 FROM pg_database WHERE datname = '{settings.DATABASE_NAME}'")
            )
            if not result.scalar():
                # 创建数据库
                conn.execute(text(f"CREATE DATABASE {settings.DATABASE_NAME}"))
                logger.info(f"✅ 数据库 '{settings.DATABASE_NAME}' 创建成功")
            else:
                logger.info(f"✅ 数据库 '{settings.DATABASE_NAME}' 已存在")
    except Exception as e:
        logger.error(f"❌ 数据库创建失败: {str(e)}")
        raise
    finally:
        postgres_engine.dispose()


def create_tables():
    """创建所有表"""
    logger.info("开始创建数据库表...")
    
    try:
        # 导入所有模型以确保它们被注册到 Base.metadata
        from app.models.user import User
        from app.models.category import Category
        from app.models.knowledge_base import KnowledgeBase
        from app.models.document import Document
        from app.models.document_share import DocumentShare
        from app.models.chat_history import ChatSession, ChatMessage
        
        logger.info("已导入所有模型")
        
        # 使用 SQLAlchemy 自动创建所有表
        Base.metadata.create_all(bind=engine)
        
        logger.info("✅ 所有表创建完成")
        
        # 额外创建一些约束（如果 ORM 没有定义）
        with engine.connect() as conn:
            # 确保 users 表的 role 约束是小写
            conn.execute(text("""
            DO $$ 
            BEGIN
                IF EXISTS (SELECT 1 FROM information_schema.table_constraints 
                          WHERE constraint_name = 'chk_role' AND table_name = 'users') THEN
                    ALTER TABLE users DROP CONSTRAINT chk_role;
                END IF;
                ALTER TABLE users ADD CONSTRAINT chk_role CHECK (role IN ('admin', 'manager', 'user', 'guest'));
            END $$;
            """))
            conn.commit()
            logger.info("✅ 用户角色约束已更新")
            
    except Exception as e:
        logger.error(f"❌ 表创建失败: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        raise


def create_initial_data():
    """创建初始数据"""
    logger.info("开始创建初始数据...")
    
    try:
        from sqlalchemy.orm import Session
        from app.models.user import User, UserRole
        from app.core.security import get_password_hash
        import uuid
        
        # 使用 ORM 创建初始数据
        with Session(engine) as session:
            # 创建默认管理员账号
            admin_user = session.query(User).filter(User.username == "admin").first()
            
            if not admin_user:
                admin_user = User(
                    id=str(uuid.uuid4()),
                    username="admin",
                    email="admin@shihe.com",
                    hashed_password=get_password_hash("admin123"),
                    full_name="系统管理员",
                    role=UserRole.ADMIN,
                    is_active=True
                )
                session.add(admin_user)
                session.commit()
                logger.info("✅ 默认管理员账号创建成功 (用户名: admin, 密码: admin123)")
            else:
                logger.info("✅ 管理员账号已存在")
            
        logger.info("✅ 初始数据创建完成")
    except Exception as e:
        logger.error(f"❌ 初始数据创建失败: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        raise


def main():
    """主函数"""
    logger.info("=" * 60)
    logger.info("🚀 开始初始化数据库...")
    logger.info("=" * 60)
    
    try:
        # 1. 检查数据库连接
        if not check_db_connection():
            logger.error("❌ 无法连接到数据库服务器")
            logger.error(f"请确保 PostgreSQL 正在运行，并检查连接信息：")
            logger.error(f"  Host: {settings.DATABASE_HOST}")
            logger.error(f"  Port: {settings.DATABASE_PORT}")
            logger.error(f"  User: {settings.DATABASE_USER}")
            sys.exit(1)
        
        # 2. 创建数据库
        create_database()
        
        # 3. 创建表
        create_tables()
        
        # 4. 创建初始数据
        create_initial_data()
        
        logger.info("=" * 60)
        logger.info("✅ 数据库初始化完成！")
        logger.info("=" * 60)
        logger.info("")
        logger.info("📝 默认管理员账号：")
        logger.info("  用户名: admin")
        logger.info("  密码: admin123")
        logger.info("")
        logger.info("⚠️  首次登录后请立即修改密码！")
        logger.info("")
        
    except Exception as e:
        logger.error("=" * 60)
        logger.error(f"❌ 数据库初始化失败: {str(e)}")
        logger.error("=" * 60)
        sys.exit(1)


if __name__ == "__main__":
    main()

